# 参数设置
d_model = 512 # 词嵌入维度d_model
n_heads = 8
d_k = d_v = d_model //n_heads   # dimension of K(=Q), V
# K,Q,V矩阵的维度（K和Q一定是一样的,因为要K乘Q的转置）,V不一定
class ScaledDotProductAttention(nn.Module):
    def __init__(self):
        super(ScaledDotProductAttention, self).__init__()
    # 添加1:加入填充遮挡attn_mask参数
    def forward(self, Q, K, V,attn_mask):
        '''
        Q: [batch_size, n_heads, len_q, d_k]
        K: [batch_size, n_heads, len_k, d_k]
        V: [batch_size, n_heads, len_v(=len_k), d_v]
        attn_mask: [batch_size, n_heads, seq_len, seq_len]
        '''

        # 所有的q和k进行矩阵相乘 (..., n, n)
        matmul_qk = torch.matmul(Q, K.transpose(-1, -2))
        d_k = torch.tensor(K.shape[-1],dtype=torch.float)
        scores = matmul_qk / np.sqrt(d_k)
        # scores : [batch_size, n_heads, len_q, len_k] (2, 8, 5, 5)

        # 添加2:scores相应位置填充为负无穷
        scores += (attn_mask * -1e9) 
        # 或：
        # scores.masked_fill_(attn_mask, -1e9)
        # 或：
        # scores=scores.masked_fill(attn_mask, -1e9)
        # masked_fill_(mask, value) 
        # 函数名后面加下划线。不经过复制操作,
        # 而是直接在原来的内存上改变它的值,可以称为原地操作符。
        # masked_fill(mask, value) 
        # 函数名后面没有下划线。经过复制操作,
        # 不是直接在原来的内存上改变它的值,而是修改复制的tensor。

        # 对每一行进行softmax,将Attention矩阵按照最后一个维度进行softmax
        attn_score = nn.Softmax(dim=-1)(scores)
        # attn_score： [batch_size, n_heads, len_q, len_k]

        # 将每一列的注意力与v进行相乘(..., n, d_v)
        context = torch.matmul(attn_score, V)
        # [batch_size, n_heads, len_q, d_v]
        # 最后返回注意力矩阵跟value的乘积,以及注意力矩阵
        return context 
